import os
import cv2

# 用户指定的文件夹路径，包含标签(txt)文件和对应的图片文件
folder_path = "UAV/labels/test"  # 替换为你的标签文件夹路径
images_folder = "UAV/images/test"  # 替换为你的图片文件夹路径

# 遍历文件夹内的标签(txt)文件
for filename in os.listdir(folder_path):
    if filename.endswith(".txt"):
        label_file_path = os.path.join(folder_path, filename)
        
        # 1. 从标签文件中读取位置信息（假设在每行的2、3、4、5位数）
        with open(label_file_path, 'r') as label_file:
            lines = label_file.readlines()
        
        # 根据用户指定的位数提取位置信息并归一化
        normalized_lines = []
        for line in lines:
            data = line.strip().split(",")
            
            # 提取2、3、4、5位作为位置信息
            x = float(data[1]) + float(data[3])/2
            y = float(data[2]) + float(data[4])/2
            width = float(data[3])
            height = float(data[4])
            
            # 获取对应的图片文件名
            image_filename = os.path.splitext(filename)[0] + ".jpg"
            image_file_path = os.path.join(images_folder, image_filename)
            
            # 获取图片的长宽
            if os.path.exists(image_file_path):
                img = cv2.imread(image_file_path)
                img_height, img_width, _ = img.shape
            else:
                img_height, img_width = 0, 0
            
            # 对位置信息进行归一化处理
            if img_width > 0 and img_height > 0:
                x_normalized = x / img_width
                y_normalized = y / img_height
                width_normalized = width / img_width
                height_normalized = height / img_height
                normalized_data = data[:1] + [str(x_normalized), str(y_normalized), str(width_normalized), str(height_normalized)] + data[5:]
                normalized_line = " ".join(normalized_data) + "\n"
                normalized_lines.append(normalized_line)
        
        # 3. 写回到标签文件
        with open(label_file_path, 'w') as label_file:
            label_file.writelines(normalized_lines)

    print(f'{filename}已完成')

print("归一化处理完成。")
